package com.jh.hts_jdk.item;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;

public class DataFrame<R, C> {
    public String name = "";
    private Object[][] data;
    private R[] row_name;
    private C[] col_name;

    private HashMap<R, Integer> row_to_index = new HashMap<>();
    private HashMap<C, Integer> col_to_index = new HashMap<>();

    public DataFrame(R[] row_name, C[] col_name, Object[][] data) {
        if (row_name.length != data.length || col_name.length != data[0].length) {
            throw new IllegalArgumentException("matrix row name or col name not match data dimension");
        }
        this.data = data;
        this.row_name = row_name.clone();
        this.col_name = col_name.clone();
        indexInit();
    }

    public DataFrame(R[] row_name, C[] col_name, Object value) {
        this(row_name, col_name, new Object[row_name.length][col_name.length]);
        for (int i = 0; i < row_name.length; i++) {
            for (int j = 0; j < col_name.length; j++) {
                this.data[i][j] = value;
            }
        }
    }

    public DataFrame(Matrix<R, C> matrix) {
        this(matrix.getRowName(), matrix.getColName(), 0);
        this.name = matrix.name;
        for (int i = 0; i < row_name.length; i++) {
            for (int j = 0; j < col_name.length; j++) {
                data[i][j] = matrix.getValue(i, j);
            }
        }
    }

    private void indexInit() {
        for (int i = 0; i < row_name.length; i++) {
            R row = row_name[i];
            if (!row_to_index.containsKey(row)) {
                row_to_index.put(row, i);
            }
        }
        for (int j = 0; j < col_name.length; j++) {
            C col = col_name[j];
            if (!col_to_index.containsKey(col)) {
                col_to_index.put(col, j);
            }
        }
    }

    public void setValue(int i, int j, Object v) {
        data[i][j] = v;
    }

    public void setValue(R row, int j, Object v) {
        setValue(row_to_index.get(row), j, v);
    }

    public void setValue(R row, C col, Object v) {
        setValue(row_to_index.get(row), col_to_index.get(col), v);
    }

    public Object getValue(int i, int j) {
        return data[i][j];
    }

    public Object getValue(R row, int j) {
        return getValue(row_to_index.get(row), j);
    }

    public Object getValue(R row, C col) {
        return getValue(row_to_index.get(row), col_to_index.get(col));
    }

    public R[] getRowName() {
        return row_name;
    }

    public C[] getColName() {
        return col_name;
    }

    public int[] getDimension() {
        return new int[] { row_name.length, col_name.length };
    }

    public void appendRow(R row_name, Object[] value) {
        if (value == null) {
            value = new Object[this.col_name.length];
        } else if (value.length != this.col_name.length) {
            value = Arrays.copyOf(value, this.col_name.length);
        }
        R[] new_row = Arrays.copyOf(this.row_name, this.row_name.length + 1);
        new_row[new_row.length - 1] = row_name;
        Object[][] new_data = new Object[new_row.length][this.col_name.length];
        for (int i = 0; i < this.row_name.length; i++) {
            for (int j = 0; j < this.col_name.length; j++) {
                new_data[i][j] = this.data[i][j];
            }
        }
        for (int j = 0; j < this.col_name.length; j++) {
            new_data[this.row_name.length][j] = value[j];
        }
        this.row_name = new_row;
        this.data = new_data;
        if (!row_to_index.containsKey(row_name)) {
            row_to_index.put(row_name, this.row_name.length - 1);
        }
    }

    public void appendCol(C col_name, Object[] value) {
        if (value == null) {
            value = new Object[this.row_name.length];
        } else if (value.length != this.row_name.length) {
            value = Arrays.copyOf(value, this.row_name.length);
        }
        C[] new_col = Arrays.copyOf(this.col_name, this.col_name.length + 1);
        new_col[new_col.length - 1] = col_name;
        Object[][] new_data = new Object[this.row_name.length][new_col.length];
        for (int i = 0; i < this.row_name.length; i++) {
            for (int j = 0; j < this.col_name.length; j++) {
                new_data[i][j] = this.data[i][j];
            }
        }
        for (int i = 0; i < this.row_name.length; i++) {
            new_data[i][this.col_name.length] = value[i];
        }
        this.col_name = new_col;
        this.data = new_data;
        if (!col_to_index.containsKey(col_name)) {
            col_to_index.put(col_name, this.col_name.length - 1);
        }
    }

    public void setRow(R row_name, Object[] v) {
        setRow(row_to_index.get(row_name), v);
    }

    public void setRow(int i, Object[] v) {
        for (int j = 0; j < col_name.length; j++) {
            if (j >= v.length) {
                break;
            }
            if (v[j] != null) {
                data[i][j] = v[j];
            }
        }
    }

    public DataFrame<C, R> T() {
        Object[][] new_data = new Object[col_name.length][row_name.length];
        for (int i = 0; i < row_name.length; i++) {
            for (int j = 0; j < col_name.length; j++) {
                new_data[j][i] = this.data[i][j];
            }
        }
        DataFrame<C, R> new_data_frame = new DataFrame<>(col_name, row_name, new_data);
        new_data_frame.name = this.name;
        return new_data_frame;
    }

    @Override
    public String toString() {
        return toString("\t",true,true);
    }

    public String toCSV(boolean show_row_name, boolean show_col_name){
        return toString(",", show_row_name, show_col_name);
    }

    public String toString(String sep, boolean show_row_name, boolean show_col_name){
        if (sep == null || sep.length() <=0) {
            sep = "\t";
        }
        StringBuilder df_str = new StringBuilder();
        List<String> row_str = new ArrayList<>();
        if (show_row_name && show_col_name) {
            row_str.add(this.name);
        }
        if (show_col_name) {
            for (C c : col_name) {
                row_str.add(c.toString());
            }
            df_str.append(String.join(sep, row_str));
            df_str.append("\n");
        }

        for (int i = 0; i < row_name.length; i++) {
            row_str.clear();
            if (show_row_name) {
                row_str.add(row_name[i].toString());
            }
            for (int j = 0; j < col_name.length; j++) {
                row_str.add(getValue(i, j).toString());
            }
            df_str.append(String.join(sep, row_str));
            df_str.append("\n");
        }

        return df_str.toString();
    }

}
